# file: custom_samplers.py

import optuna

class GridSamplerWrapper(optuna.samplers.GridSampler):
    """
    A wrapper for GridSampler that accepts and ignores a 'seed' argument.
    
    This prevents a TypeError when Hydra automatically injects a seed into
    a sampler that doesn't support it.
    """
    def __init__(self, *args, **kwargs):
        # Remove 'seed' from the keyword arguments if it exists., **kwargs
        # print(type(kwargs))
        # print(kwargs.keys())
        # for k in ['consider_prior', 'seed', 'prior_weight', 'consider_magic_clip']:
        #     kwargs.pop(k, None)
        parameter_name = "model.lr_scales.4"

        # The choices from your YAML file as a Python list
        parameter_choices = [1, 0.5, 1e-1]

        # 'search_space' is a dictionary containing the name and choices
        search_space = {
            parameter_name: parameter_choices
        }

        # Call the original GridSampler's constructor with the cleaned arguments.
        super().__init__(search_space=search_space, *args)